using BootstrapBlazorApp.Server.LLMs.api;
using BootstrapBlazorApp.Server.LLMs.messages;

namespace BootstrapBlazorApp.Server.LLMs.chain;
//将各个组件组织成链
public class SequenceChain : IRunnable
{
    private readonly List<IRunnable> _nodes;

    private SequenceChain(SequenceChain chain, IRunnable node)
    {
        _nodes = chain._nodes;
        _nodes.Add(node);
    }

    private SequenceChain(IRunnable node)
    {
        _nodes = [node];
    }

    public async Task<ChainContext> InvokeAsync(ChainContext input)
    {
        for (int i = 0; i < _nodes.Count - 1; i++)
        {
            await _nodes[i].InvokeAsync(input);
        }

        return await _nodes.Last().InvokeAsync(input);
    }

    public async IAsyncEnumerable<ChainContext> InvokeStreamAsync(ChainContext input)
    {
        if (_nodes.Count == 1)
        {
            await foreach (var context in _nodes[0].InvokeStreamAsync(input))
            {
                yield return context;
            }
        }
        else
        {
            for (var i = 0; i < _nodes.Count-2; i++)
            {
                await _nodes[i].InvokeAsync(input);
            }
            //流式返回最后添加的输出解析器的结果
            await foreach (var context in _nodes[^2].InvokeStreamAsync(input))
            {
                await foreach (var chainContext in _nodes[^1].InvokeStreamAsync(context))
                {
                    yield return chainContext;
                }

            }
        }
        
    }

    public async Task<ChainContext> InvokeAsync(string s)
    {
        var input = BuildContext(s);
        return await InvokeAsync(input);
    }
    public async Task<ChainContext> InvokeAsync(Dictionary<string,string> dictionary)
    {
        var input = BuildContext(dictionary);
        return await InvokeAsync(input);
    }
    public IAsyncEnumerable<ChainContext> InvokeStreamAsync(Dictionary<string,string> dictionary)
    {
        var input = BuildContext(dictionary);
        return InvokeStreamAsync(input);
    }

    public IAsyncEnumerable<ChainContext> InvokeStreamAsync(string s)
    {
        var input = BuildContext(s);
        return InvokeStreamAsync(input);
    }

    public SequenceChain Pipeline(IRunnable node)
    {
        return new SequenceChain(this, node);
    }

    public static SequenceChain Init(IRunnable node)
    {
        return new SequenceChain(node);
    }

    private ChainContext BuildContext(string s)
    {
        ChainContext input = new ChainContext();
        input.Add(new HumanMessage(s));
        foreach (var node in _nodes)
        {
            if (node is IBaseOutPutParser parser)
            {
                input.Add(new SystemMessage(parser.GetInstruction()));
            }
        }
        return input;
    }
    
    private ChainContext BuildContext(Dictionary<string,string> dictionary)
    {
        ChainContext input = new ChainContext { PromptInputValue = dictionary };
        foreach (var node in _nodes)
        {
            if (node is IBaseOutPutParser parser)
            {
                input.Add(new SystemMessage(parser.GetInstruction()));
            }
        }
        return input;
    }
    
}
